import torch.nn.utils.rnn as rnn
from torch.utils.data import DataLoader
from dataset import *
import parameter_config

params = parameter_config.ParameterConfig()

def get_dataset(data_path):
    with open(data_path, 'rb') as f:
        data = pickle.load(f)
    my_dataset = MyDataset(data, max_len=params.max_len)
    return my_dataset

def collate_fn(batch):
    # 利用rnn工具进行填充，使得其输入序列的长度一致
    input_ids = rnn.pad_sequence(batch, batch_first=True, padding_value=0)
    labels = rnn.pad_sequence(batch, batch_first=True, padding_value=-100)
    return input_ids, labels

def get_dataloader(train_pkl_path, valid_pkl_path):
    """
    获取训练集和验证机的dataloader对象
    :param train_pkl_path:训练集路径
    :param valid_pkl_path: 验证机路径
    :return: 训练集和验证机的dataloader对象
    """
    train_dataset = get_dataset(train_pkl_path)
    val_dataset = get_dataset(valid_pkl_path)
    train_dataloader = DataLoader(train_dataset,
                                  batch_size= params.batch_size,
                                  shuffle=True,
                                  collate_fn=collate_fn,
                                  drop_last=True)
    val_dataloader = DataLoader(val_dataset,
                                batch_size= params.batch_size,
                                shuffle=True,
                                collate_fn=collate_fn,
                                drop_last=True)
    return train_dataloader, val_dataloader
